{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 实体检索树（TrieTree）"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TrieTree 存储字符串、检索   【同名实体 '_'后字符作为列表存储】\n",
    "class TrieTree(object):\n",
    "    \"\"\"\n",
    "    字典树用于字符串的存储、检索\n",
    "    【.】特殊符号：代表末端节点（字符串中不能出现此字符）\n",
    "    【_】特殊符号：代表有多个同名实体,\"_\"符号后面内容为实体备注信息\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self):\n",
    "        self.tree = {}  # 字典树\n",
    "        self.deep = 0  # TrieTree 最大深度\n",
    "        self.node_num = 0  # TrieTree 存储字符串数量\n",
    "\n",
    "        self.values_list = list()  # TrieTree转化成字符串\n",
    "\n",
    "        self.match_word = list()  # 待检索的句子\n",
    "        self.match_entities = list()  # 检索到的实体\n",
    "\n",
    "    # TrieTree 加入节点\n",
    "    def add(self, word):\n",
    "        \"\"\"\n",
    "        :param word:  待加入字符串\n",
    "        :return:\n",
    "        \"\"\"\n",
    "        tree = self.tree\n",
    "\n",
    "        multi_flag = False  # 是否多同名实体 （是否有 '_' 特殊符号）\n",
    "        for i, char in enumerate(word):\n",
    "\n",
    "            if '_' == char:\n",
    "                multi_flag = True\n",
    "\n",
    "                if '_' not in tree:\n",
    "                    tree['_'] = [word[i + 1:]]\n",
    "                else:\n",
    "                    tree['_'].append(word[i + 1:])\n",
    "\n",
    "                break\n",
    "            elif char in tree:\n",
    "                tree = tree[char]\n",
    "\n",
    "            else:\n",
    "                tree[char] = {}\n",
    "                tree = tree[char]\n",
    "\n",
    "        if not multi_flag:\n",
    "            tree['.'] = {}  # 【.】特殊符号，代表末端节点（字符串中不能出现此字符）\n",
    "\n",
    "    # TrieTree 检索  全匹配模式\n",
    "    def entity_search_f(self, search_str):\n",
    "        self.match_word = list()  # 清空 待检索的句子\n",
    "        self.match_entities = list()  # 清空 检索到的实体\n",
    "\n",
    "        # 全匹配模式检索\n",
    "        self._entity_search_f(search_str)\n",
    "\n",
    "        # print(self.match_word)\n",
    "\n",
    "        # 返回检索到的实体列表\n",
    "        return self.match_entities\n",
    "\n",
    "    def _entity_search_f(self, word):\n",
    "        \"\"\"\n",
    "        :param word:  待检索字符串\n",
    "        :return:\n",
    "        \"\"\"\n",
    "        # 当前句子未进行检索\n",
    "        if word not in self.match_word:\n",
    "            self.match_word.append(word)\n",
    "\n",
    "            # print(word)\n",
    "            tree = self.tree\n",
    "\n",
    "            for i, char in enumerate(word):\n",
    "                if char in tree:\n",
    "                    tree = tree[char]\n",
    "\n",
    "                    # 实体结束标志\n",
    "                    if '.' in tree:\n",
    "                        # 保存检索到实体名称\n",
    "                        self.match_entities.append(word[:i + 1])\n",
    "\n",
    "                        # 未结束，继续检索\n",
    "                        if len(word) > i + 1:\n",
    "                            # print('N ', end='')\n",
    "                            self._entity_search_f(word[i + 1:])\n",
    "\n",
    "                    # 是否为多同名实体 （是否有 '_' 特殊符号）\n",
    "                    if '_' in tree:\n",
    "                        # 保存检索到的所有同名实体名称\n",
    "                        match_result = [word[: i + 1] + j for j in self.get_values_list(tree) if j.startswith('_')]\n",
    "                        self.match_entities.extend(match_result)\n",
    "\n",
    "                        # 未结束，继续检索\n",
    "                        if len(word) > i + 1:\n",
    "                            # print('N ', end='')\n",
    "                            self._entity_search_f(word[i + 1:])\n",
    "\n",
    "                else:\n",
    "\n",
    "                    # 未结束，跳过一个字符，继续检索\n",
    "                    if len(word) > i + 1:\n",
    "                        # print('E ', end='')\n",
    "                        self._entity_search_f(word[i + 1:])\n",
    "\n",
    "                    break\n",
    "\n",
    "    # TrieTree 检索  最大匹配模式 贪心模式\n",
    "    def entity_search_g(self, search_str):\n",
    "        index = 0\n",
    "\n",
    "        match_entities = list()\n",
    "        for _ in range(999):  # 检索句子最长999\n",
    "\n",
    "            # 检索实体\n",
    "            search_result = self._entity_search_g(search_str[index:])\n",
    "\n",
    "            # 检索到实体\n",
    "            if len(search_result) > 0:\n",
    "                # print(search_result)  # 检索到的实体名称\n",
    "\n",
    "                # 根据限制条件保留检索到实体\n",
    "                for s_i in search_result:\n",
    "                    match_entities.append(s_i)\n",
    "\n",
    "                if len(search_result) == 1:\n",
    "                    index += len(search_result[0])  # 只检索到一个实体\n",
    "                else:\n",
    "                    index += len(search_result[0].split('_')[0])  # 检索到多个实体，只实体'_'分隔符前\n",
    "\n",
    "            # 检索结束\n",
    "            elif index >= len(search_str):\n",
    "                # print('Search Finish.')\n",
    "                break\n",
    "\n",
    "            # 检索不到对应实体\n",
    "            else:\n",
    "                # print('UnSearch:【{}】'.format(search_str[index]))\n",
    "                index += 1  # 跳过一个字符 继续检索\n",
    "\n",
    "        return match_entities\n",
    "\n",
    "    def _entity_search_g(self, word):\n",
    "        \"\"\"\n",
    "        :param word:  待检索字符串\n",
    "        :return:\n",
    "        \"\"\"\n",
    "        tree = self.tree\n",
    "\n",
    "        max_index = 0  # 最长匹配长度\n",
    "        multi_flag = False  # 是否多同名实体 （是否有 '_' 特殊符号）\n",
    "\n",
    "        for i, char in enumerate(word):\n",
    "            if char in tree:\n",
    "                tree = tree[char]\n",
    "\n",
    "                # 实体结束标志\n",
    "                if '.' in tree:\n",
    "                    max_index = max(max_index, i + 1)\n",
    "                    multi_flag = False\n",
    "\n",
    "                # 多同名实体标志\n",
    "                if '_' in tree:\n",
    "                    max_index = max(max_index, i + 1)\n",
    "                    multi_flag = True\n",
    "\n",
    "            else:\n",
    "                # 字符串已到达最大长度\n",
    "                break\n",
    "\n",
    "        # 检索结果为同名实体，提取所有同名实体名称\n",
    "        if multi_flag:\n",
    "            tree = self.tree\n",
    "            for word_i in word[: max_index]:\n",
    "                tree = tree[word_i]\n",
    "\n",
    "            # 所有同名实体名称\n",
    "            match_result = [word[: max_index] + i for i in self.get_values_list(tree) if i.startswith('_')]\n",
    "            # 是否有直接结束实体 TODO 图谱存储结构不完善，同名实体应都以 '_' 加备注\n",
    "            if '.' in tree:\n",
    "                match_result.append(word[: max_index])\n",
    "\n",
    "            return match_result\n",
    "\n",
    "        # 未匹配到任何信息\n",
    "        if max_index == 0:\n",
    "            return []\n",
    "        # 匹配到单实体信息\n",
    "        else:\n",
    "            return [word[: max_index]]\n",
    "\n",
    "    # TrieTree信息（深度、节点数量）\n",
    "    def info(self):\n",
    "        # TrieTree 存储字符串数量\n",
    "        self.node_num = self._info(self.tree)\n",
    "\n",
    "        # self.deep 包含结束节点层 .  比实际字符串长度+1\n",
    "        return {'Deep': self.deep, 'Num': self.node_num}\n",
    "\n",
    "    def _info(self, t, deep=0, num=0):\n",
    "        self.deep = max(self.deep, deep)\n",
    "        # print('| ' * deep, t, num)\n",
    "\n",
    "        if len(t) > 0:\n",
    "            for t_i in t:\n",
    "                if t_i != '_':\n",
    "                    num = self._info(t[t_i], deep + 1, num)\n",
    "                else:\n",
    "                    num = num + len(t[t_i])\n",
    "        else:\n",
    "            num = num + 1\n",
    "\n",
    "        return num\n",
    "\n",
    "    # 将trie_tree 转化成字符串\n",
    "    def get_values_list(self, tree_x):\n",
    "        self.values_list = list()  # 初始化\n",
    "        self._get_values_list(tree_x)\n",
    "        return self.values_list  # 返回tree中所有字符串\n",
    "\n",
    "    def _get_values_list(self, trees, value=''):\n",
    "        for c_i in trees:\n",
    "\n",
    "            # 结束\n",
    "            if c_i == '.':\n",
    "                self.values_list.append(value)\n",
    "                continue\n",
    "\n",
    "            # 多实体\n",
    "            if c_i == '_':\n",
    "                # 结束，已到达末端节点\n",
    "                self.values_list.extend([value + '_' + t_i for t_i in trees[c_i]])\n",
    "                continue\n",
    "\n",
    "            # 未结束\n",
    "            if trees[c_i] != {'.': {}}:\n",
    "                self._get_values_list(trees[c_i], value + c_i)\n",
    "\n",
    "            else:\n",
    "                # 结束，已到达末端节点\n",
    "                self.values_list.append(value + c_i)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "字典树大小:   {'Deep': 5, 'Num': 7}\n",
      "存储节点信息:   ['新加坡', '新加坡_（国家）', '新加坡_（城市）', '新加坡2', '新加', '新', '水域率']\n",
      "全匹配模式:  ['新', '水域率', '新加', '新加坡', '新加坡_（国家）', '新加坡_（城市）']\n",
      "最大匹配模式:  ['新加坡_（国家）', '新加坡_（城市）', '新加坡', '水域率']\n"
     ]
    }
   ],
   "source": [
    "trie_tree = TrieTree()\n",
    "\n",
    "# 实体列表 \n",
    "entities = ['新加坡','新加坡_（国家）','新加坡_（城市）','新加坡2','新','新加','水域率']\n",
    "\n",
    "for entity_i in entities:\n",
    "    trie_tree.add(entity_i)\n",
    "    \n",
    "print('字典树大小:  ',trie_tree.info() )\n",
    "# 存储节点信息\n",
    "print('存储节点信息:  ', trie_tree.get_values_list(trie_tree.tree))\n",
    "\n",
    "\n",
    "search_str = '新加坡水域率'\n",
    "# 节点检索  全匹配模式\n",
    "print('全匹配模式: ', trie_tree.entity_search_f(search_str))\n",
    "# 节点检索  最大匹配模式\n",
    "print('最大匹配模式: ', trie_tree.entity_search_g(search_str))\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "##  准确的获取一个对象（比如字典）的内存占用"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "848\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "240 <class 'dict'> {'a': 1, 'b': 2, 'c': 3, 'd': [4, 5, 6, 7], ...}\n",
      "50 <class 'str'> 'a'\n",
      "28 <class 'int'> 1\n",
      "50 <class 'str'> 'b'\n",
      "28 <class 'int'> 2\n",
      "50 <class 'str'> 'c'\n",
      "28 <class 'int'> 3\n",
      "50 <class 'str'> 'd'\n",
      "96 <class 'list'> [4, 5, 6, 7]\n",
      "28 <class 'int'> 4\n",
      "28 <class 'int'> 5\n",
      "28 <class 'int'> 6\n",
      "28 <class 'int'> 7\n",
      "50 <class 'str'> 'e'\n",
      "66 <class 'str'> 'a string of chars'\n"
     ]
    }
   ],
   "source": [
    "from __future__ import print_function\n",
    "from sys import getsizeof, stderr\n",
    "from itertools import chain\n",
    "from collections import deque\n",
    "try:\n",
    "    from reprlib import repr\n",
    "except ImportError:\n",
    "    pass\n",
    "\n",
    "def total_size(o, handlers={}, verbose=False):\n",
    "    \"\"\" Returns the approximate memory footprint an object and all of its contents.\n",
    "\n",
    "    Automatically finds the contents of the following builtin containers and\n",
    "    their subclasses:  tuple, list, deque, dict, set and frozenset.\n",
    "    To search other containers, add handlers to iterate over their contents:\n",
    "\n",
    "        handlers = {SomeContainerClass: iter,\n",
    "                    OtherContainerClass: OtherContainerClass.get_elements}\n",
    "\n",
    "    \"\"\"\n",
    "    dict_handler = lambda d: chain.from_iterable(d.items())\n",
    "    all_handlers = {tuple: iter,\n",
    "                    list: iter,\n",
    "                    deque: iter,\n",
    "                    dict: dict_handler,\n",
    "                    set: iter,\n",
    "                    frozenset: iter,\n",
    "                   }\n",
    "    all_handlers.update(handlers)     # user handlers take precedence\n",
    "    seen = set()                      # track which object id's have already been seen\n",
    "    default_size = getsizeof(0)       # estimate sizeof object without __sizeof__\n",
    "\n",
    "    def sizeof(o):\n",
    "        if id(o) in seen:       # do not double count the same object\n",
    "            return 0\n",
    "        seen.add(id(o))\n",
    "        s = getsizeof(o, default_size)\n",
    "\n",
    "        if verbose:\n",
    "            print(s, type(o), repr(o), file=stderr)\n",
    "\n",
    "        for typ, handler in all_handlers.items():\n",
    "            if isinstance(o, typ):\n",
    "                s += sum(map(sizeof, handler(o)))\n",
    "                break\n",
    "        return s\n",
    "\n",
    "    return sizeof(o)\n",
    "\n",
    "\n",
    "##### Example call #####\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    d = dict(a=1, b=2, c=3, d=[4,5,6,7], e='a string of chars')\n",
    "    print(total_size(d, verbose=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "240 <class 'dict'> {'新': {'.': {}, '加': {'.': {}, '坡': {'.': {}, '2': {'.': {}}, '_': ['（国家）', '（城市）']}}}, '水': {'域': {'率': {'.': {}}}}}\n",
      "76 <class 'str'> '新'\n",
      "240 <class 'dict'> {'.': {}, '加': {'.': {}, '坡': {'.': {}, '2': {'.': {}}, '_': ['（国家）', '（城市）']}}}\n",
      "76 <class 'str'> '加'\n",
      "240 <class 'dict'> {'.': {}, '坡': {'.': {}, '2': {'.': {}}, '_': ['（国家）', '（城市）']}}\n",
      "76 <class 'str'> '坡'\n",
      "240 <class 'dict'> {'.': {}, '2': {'.': {}}, '_': ['（国家）', '（城市）']}\n",
      "50 <class 'str'> '.'\n",
      "240 <class 'dict'> {}\n",
      "50 <class 'str'> '_'\n",
      "104 <class 'list'> ['（国家）', '（城市）']\n",
      "82 <class 'str'> '（国家）'\n",
      "82 <class 'str'> '（城市）'\n",
      "50 <class 'str'> '2'\n",
      "240 <class 'dict'> {'.': {}}\n",
      "240 <class 'dict'> {}\n",
      "240 <class 'dict'> {}\n",
      "240 <class 'dict'> {}\n",
      "76 <class 'str'> '水'\n",
      "240 <class 'dict'> {'域': {'率': {'.': {}}}}\n",
      "76 <class 'str'> '域'\n",
      "240 <class 'dict'> {'率': {'.': {}}}\n",
      "76 <class 'str'> '率'\n",
      "240 <class 'dict'> {'.': {}}\n",
      "240 <class 'dict'> {}\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "3994"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "total_size(trie_tree.tree, verbose=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
